import numpy as np
import networkx as nx
from src import simulation
import random
import torch
from scipy.sparse import coo_matrix
import scipy.sparse as sp

class Enviroment:
    def __init__(
            self,
            max_add_edges_num = 10,
            seed_num = 5,
            candidate_edges_num = 50
                 ):
        self.max_add_edges_num = max_add_edges_num
        self.seed_num = seed_num
        self.candidate_edges_num = candidate_edges_num

    def reset(self):
        self.G = nx.random_graphs.barabasi_albert_graph(50,2)
        self.node_prob = {}
        self.edge_prob = {}
        self.current_add_edges_num = 0
        self.supplement_edges = []
        #candidate info
        self.candidate_edges = []
        self.candidate_weights = []
        self.candidate_mask = np.ones(self.candidate_edges_num)
        self.seeds = np.zeros(self.G.number_of_nodes())

        #transfer G to digraph
        g0 = nx.random_graphs.barabasi_albert_graph(50, 2)
        edgelist = [list(e) for e in g0.edges()]
        new_edgelist = []
        for edge in edgelist:
            event = np.random.random_sample()
            if event > 0.5:
                edge.reverse()
                new_edgelist.append(edge)
            else:
                new_edgelist.append(edge)

        self.G = nx.DiGraph()
        self.G.add_edges_from(new_edgelist)


        #get node prob
        for i in self.G.nodes:
            self.node_prob[i] = 1

        #get supplement_edges
        for i in self.G.nodes:
            for j in self.G.nodes:
                if not self.G.has_edge(i,j) and i != j:
                    self.supplement_edges.append((i,j))

        #get candidate_edges and mask
        self.candidate_edges = random.sample(self.supplement_edges,self.candidate_edges_num)
        self.candidate_weights = [1 / (len(list(self.G.predecessors(list(e)[-1]))) + 1) for e in self.candidate_edges]
        self.tensor_candidate_edges = torch.Tensor(self.candidate_edges).t().long() #[2,C]
        self.tensor_candidate_mask = torch.Tensor(self.candidate_mask).unsqueeze(0) #[1,C]
        self.tensor_candidate_weights = torch.Tensor(self.candidate_weights).unsqueeze(0).t() #[C, 1]

        #candidate_tensor = torch.cat((self.candidate_edges,self.candidate_mask),0)

        #sample seeds
        idx = random.sample(range(self.G.number_of_nodes()),self.seed_num)
        self.seeds[idx] = 1
        self.tensor_seeds = torch.Tensor(self.seeds).reshape(-1, 1)
        self.seeds_list = idx

        #set graph weight
        self.edge_prob = {e: 1 / len(list(self.G.predecessors(list(e)[-1]))) for e in self.G.edges()}
        nx.set_edge_attributes(self.G, self.edge_prob, "weight")
        weight_adj = nx.adj_matrix(self.G).toarray()
        adj = torch.Tensor(weight_adj)

        #print(adj,self.seeds)
        state = [adj, self.tensor_seeds, self.tensor_candidate_edges, self.tensor_candidate_weights, self.tensor_candidate_mask]

        self.spreading = simulation.SIMU("SIR",self.G,self.node_prob, self.edge_prob, self.seeds_list, path = "", num_simulation=10000, cpu_core=20)[-1, :, -1].sum()
        return state

    def step(self, action):
        #update graph
        self.candidate_mask[action] = 0
        u,v = self.candidate_edges[action]
        self.G.add_edge(u,v)
        self.G.edges[(u, v)]["weight"] = 1 / (len(list(self.G.predecessors(v))) + 1)
        self.tensor_candidate_mask = torch.Tensor(self.candidate_mask).unsqueeze(0)  # [1,C]

        #get reward
        spreading = simulation.SIMU("SIR",self.G,self.node_prob, self.edge_prob, self.seeds_list, path = "", num_simulation=10000, cpu_core=20)[-1, :, -1].sum()

        reward = spreading - self.spreading
        self.spreading = spreading
        #get state
        weight_adj = nx.adj_matrix(self.G).toarray()
        adj = torch.Tensor(weight_adj)
        state = [adj, self.tensor_seeds, self.tensor_candidate_edges, self.tensor_candidate_weights, self.tensor_candidate_mask]

        #check if has done
        done = False
        if self.current_add_edges_num + 1 ==  self.max_add_edges_num:
            done = True
        self.current_add_edges_num += True


        return state, reward, done


if __name__ == '_main_':
    env = Enviroment(max_add_edges_num = 10,seed_num = 5)
    state = env.reset()
    env.step(1)
